package net.zjvis.lab.nebula.backend.lib.algorithm.conf.algorithm;

import com.alibaba.fastjson.JSONObject;
import java.util.HashMap;
import java.util.Map;
import org.zjvis.datascience.common.enums.AlgPyEnum;

public class AlgorithmHelper {

    public static final Map<String, String> algorithmNameMap = new HashMap<>();

    static {
        algorithmNameMap.put("new_kmeans", "KMeans");
        algorithmNameMap.put("new_dbscan", "DBscan");
        algorithmNameMap.put("new_pca", "Pca");
        algorithmNameMap.put("new_tsne", "Tsne");
        algorithmNameMap.put("new_lle", "LLE");
        algorithmNameMap.put("new_iso_forest", "IsoForest");
        algorithmNameMap.put("new_fp_growth", "FPGrowth");
        algorithmNameMap.put("new_prefix_span", "PrefixSpan");
        algorithmNameMap.put("new_arima", "ARIMA");
        algorithmNameMap.put("new_holt_winters", "Holt_Winters");
        algorithmNameMap.put("new_shift", "Shift");
        algorithmNameMap.put("new_adf", "ADF");
    }



    public static Algorithm parseFromJsonObject(JSONObject jsonObject) {
        AlgPyEnum algPyEnum = AlgPyEnum.getEnumByVal(jsonObject.getInteger("algType"));
        JSONObject formData = jsonObject.getJSONArray("setParams").getJSONObject(0).getJSONObject("formData");
        if (algPyEnum.equals(AlgPyEnum.DBSCAN)) {
            return DBSCAN.builder().build();
        } else if (algPyEnum.equals(AlgPyEnum.KMEANS)) {
            return KMeans.builder()
                    .clusterCount(formData.getInteger("k"))
                    .maxIteration(formData.getInteger("max_iter"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.KMEANS_NEW)) {
            return KMeans.builder()
                    .clusterCount(formData.getInteger("k"))
                    .maxIteration(formData.getInteger("max_iter"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.DBSCAN_NEW)) {
            return DBSCAN.builder()
                    .minSamples(formData.getInteger("min_samples"))
                    .metric(formData.getString("metric"))
                    .eps(formData.getDouble("eps"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.PCA_NEW)) {
            return PCA.builder()
                    .k(formData.getInteger("k"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.TSNE_NEW)) {
            return TSNE.builder()
                    .k(formData.getInteger("k"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.LLE_NEW)) {
            return LLE.builder()
                    .components(formData.getInteger("n_components"))
                    .neighbors(formData.getInteger("n_neighbors"))
                    .maxIteration(formData.getInteger("max_iter"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.ISO_FOREST_NEW)) {
            return IsoForest.builder()
                    .estimators(formData.getInteger("n_estimators"))
                    .contamination(formData.getDouble("contamination"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.FP_GROWTH_NEW)) {
            return FpGrowth.builder()
                    .minSupport(formData.getDouble("min_support"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.PREFIX_SPAN_NEW)) {
            return PrefixSpan.builder()
                    .minSupport(formData.getDouble("min_support"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.ARIMA_NEW)) {
            return ARIMA.builder()
                    .col(formData.getJSONArray("cols").getString(0))
                    .q(formData.getInteger("q"))
                    .d(formData.getInteger("d"))
                    .p(formData.getInteger("p"))
                    .steps(formData.getInteger("steps"))
                    .trend(formData.getString("trend"))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.HOLT_WINTERS_NEW)) {
            return HoltWinters.builder()
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .steps(formData.getInteger("steps"))
                    .period(formData.getInteger("period"))
                    .trend(formData.getString("trend"))
                    .seasonal(formData.getString("seasonal"))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.SHIFT_NEW)) {
            return Shift.builder()
                    .in(formData.getInteger("n_in"))
                    .out(formData.getInteger("n_out"))
                    .drop(formData.getBoolean("drop"))
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.ADF_NEW)) {
            return ADF.builder()
                    .columnList(formData.getJSONArray("cols").toJavaList(String.class))
                    .autolag(formData.getString("autolag"))
                    .pvalue(formData.getDouble("pvalue"))
                    .maxlag(formData.getInteger("maxlag"))
                    .regresssion(formData.getString("regression"))
                    .build();
        } else if (algPyEnum.equals(AlgPyEnum.FEATURE_SCALING)) {
            return Scaling.builder().build();
        }
        return null;
    }
}
